package com.cetc.sdp.kmga.cs.util;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.*;
import org.apache.hadoop.hbase.mapreduce.TableOutputFormat;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @description:
 * @author： DengQiang
 * @date: 2017/10/25 13:42
 */
public class PooledHbaseTableOutputFormat<KEY> extends TableOutputFormat {

    private static final Log LOG = LogFactory.getLog(PooledHbaseTableOutputFormat.class);

    private static ConcurrentHashMap<String, Map<String, PooledHbaseTableOutputFormat.CustomerTableRecordWriter>> writerConcurrentHashMap = new ConcurrentHashMap<>(32);

    private class CustomerTableRecordWriter
            extends RecordWriter<KEY, Mutation> {

        private Connection connection;
        private BufferedMutator mutator;

        private Configuration conf = null;

        /**
         * @throws IOException
         */
        public CustomerTableRecordWriter() throws IOException {
            conf = getConf();
            String tableName = conf.get(OUTPUT_TABLE);
            this.connection = ConnectionFactory.createConnection(conf);
            this.mutator = connection.getBufferedMutator(TableName.valueOf(tableName));
            LOG.info("Created table instance for " + tableName);
        }

        /**
         * Closes the writer, in this case flush table commits.
         *
         * @param context The context.
         * @throws IOException When closing the writer fails.
         * @see RecordWriter#close(TaskAttemptContext)
         */
        @Override
        public void close(TaskAttemptContext context)
                throws IOException {
            mutator.flush();
        }

        /**
         * Writes a key/value pair into the table.
         *
         * @param key   The key.
         * @param value The value.
         * @throws IOException When writing fails.
         * @see RecordWriter#write(Object, Object)
         */
        @Override
        public void write(KEY key, Mutation value)
                throws IOException {
            if (!(value instanceof Put) && !(value instanceof Delete)) {
                throw new IOException("Pass a Delete or a Put");
            }
            mutator.mutate(value);
        }
    }

    @Override
    public RecordWriter getRecordWriter(TaskAttemptContext context) throws IOException, InterruptedException {
        Map<String, PooledHbaseTableOutputFormat.CustomerTableRecordWriter> map = writerConcurrentHashMap
                .computeIfAbsent(Thread.currentThread().getName(), k -> new HashMap<>(16));
        String key = getConf().get(OUTPUT_TABLE);
        PooledHbaseTableOutputFormat.CustomerTableRecordWriter writer = map.get(key);
        if (writer == null) {
            System.out.println("Get new CustomerTableRecordWriter for Thread:" + Thread.currentThread().getName());
            writer = new CustomerTableRecordWriter();
            map.put(key, writer);
        }
        if (writer.connection.isClosed()) {
            writer = new CustomerTableRecordWriter();
            map.put(key, writer);
        }
        return writer;
    }
}
